{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Build a ChatBot "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have built a simple question-answering pipeline, where we can input a question and get an answer. In addition to single round of question-answering, we can also have a conversation with an LLM by building a chatbot. The chatbot can remember the history of the conversation and respond based on the history. The key to achieve this is to leverage the promt args `chat_history_str` and the data structure `Memory` to manage the conversation history."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import needed modules from LightRAG\n",
    "from lightrag.core.component import Component\n",
    "from lightrag.core.generator import Generator\n",
    "from lightrag.core.memory import Memory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Here, we use the OpenAIClient as an example, but you can use any other clients (with the corresponding API Key as needed), such as AnthropicAPIClient\n",
    "from lightrag.components.model_client import OpenAIClient\n",
    "OPENAI_API_KEY=\"YOUR_API_KEY\" # Replace with your OpenAI API Key, or you can put it in a .env file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ChatBot(\n",
      "  (generator): Generator(\n",
      "    model_kwargs={'model': 'gpt-3.5-turbo'}, model_type=ModelType.LLM\n",
      "    (system_prompt): Prompt(\n",
      "      template: {# task desc #}\n",
      "      {% if task_desc_str %}\n",
      "      {{task_desc_str}}\n",
      "      {% else %}\n",
      "      Answer user query.\n",
      "      {% endif %}\n",
      "      {# output format #}\n",
      "      {% if output_format_str %}\n",
      "      <OUTPUT_FORMAT>\n",
      "      {{output_format_str}}\n",
      "      </OUTPUT_FORMAT>\n",
      "      {% endif %}\n",
      "      {# tools #}\n",
      "      {% if tools_str %}\n",
      "      <TOOLS>\n",
      "      {{tools_str}}\n",
      "      </TOOLS>\n",
      "      {% endif %}\n",
      "      {# example #}\n",
      "      {% if examples_str %}\n",
      "      <EXAMPLES>\n",
      "      {{examples_str}}\n",
      "      </EXAMPLES>\n",
      "      {% endif %}\n",
      "      {# chat history #}\n",
      "      {% if chat_history_str %}\n",
      "      <CHAT_HISTORY>\n",
      "      {{chat_history_str}}\n",
      "      </CHAT_HISTORY>\n",
      "      {% endif %}\n",
      "      {#contex#}\n",
      "      {% if context_str %}\n",
      "      <CONTEXT>\n",
      "      {{context_str}}\n",
      "      </CONTEXT>\n",
      "      {% endif %}\n",
      "      {# steps #}\n",
      "      {% if steps_str %}\n",
      "      <STEPS>\n",
      "      {{steps_str}}\n",
      "      </STEPS>\n",
      "      {% endif %}\n",
      "      {% if input_str %}\n",
      "      <Inputs>\n",
      "      {{input_str}}\n",
      "      </Inputs>\n",
      "      {% endif %}\n",
      "      {% if output_str %}\n",
      "      <Outputs>\n",
      "      {{output_str}}\n",
      "      </Outputs>\n",
      "      {% endif %}\n",
      "      , prompt_variables: ['context_str', 'task_desc_str', 'tools_str', 'chat_history_str', 'input_str', 'output_str', 'output_format_str', 'steps_str', 'examples_str']\n",
      "    )\n",
      "    (model_client): OpenAIClient()\n",
      "  )\n",
      "  (chat_history): Memory()\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "# Build the ChatBot pipeline\n",
    "class ChatBot(Component):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.generator = Generator(\n",
    "            model_client=OpenAIClient(),\n",
    "            model_kwargs={'model': 'gpt-3.5-turbo'}\n",
    "        )\n",
    "        self.chat_history = Memory() # Memory to store the chat history\n",
    "    \n",
    "    def call(self) -> str:\n",
    "        print(\"Welcome to the ChatBot. Type anything to chat. Type 'exit' to end.\")\n",
    "        while True:\n",
    "            user_input = input(\"You: \")\n",
    "            if user_input.lower() == \"exit\":\n",
    "                print(\"Goodbye!\")\n",
    "                break\n",
    "            chat_history_str = self.chat_history()\n",
    "            # Generate the response from the user input and chat history\n",
    "            response = self.generator(\n",
    "                prompt_kwargs={\n",
    "                    \"input_str\": user_input,\n",
    "                    \"chat_history_str\": chat_history_str,\n",
    "                },\n",
    "            )\n",
    "            # Save the user input and response to the memory\n",
    "            self.chat_history.add_dialog_turn(\n",
    "                user_query=user_input, assistant_response=response\n",
    "            )\n",
    "            print(f\"ChatBot: {response}\")\n",
    "\n",
    "chatbot = ChatBot()\n",
    "print(chatbot)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Welcome to the ChatBot. Type anything to chat. Type 'exit' to end.\n",
      "ChatBot: GeneratorOutput(data=\"Learning to drive can be an exciting and rewarding experience. Here are some general steps to help you get started on your journey to becoming a safe and confident driver:\\n\\n1. Get a learner's permit: In most places, you will need to obtain a learner's permit before you can start learning how to drive. Check with your local department of motor vehicles for the specific requirements in your area.\\n\\n2. Take a driver's education course: Consider enrolling in a driver's education course to learn the rules of the road and get some hands-on practice with a qualified instructor.\\n\\n3. Practice with a licensed driver: Before you can get your driver's license, you will need to log a certain number of supervised driving hours with a licensed adult. This is a great opportunity to get comfortable behind the wheel and practice your skills.\\n\\n4. Study the driver's manual: Make sure to familiarize yourself with the driver's manual for your state or country. It contains important information about traffic laws, road signs, and safe driving practices.\\n\\n5. Practice, practice, practice: The more time you spend behind the wheel, the more confident you will become as a driver. Practice in a variety of different conditions and situations to hone your skills.\\n\\n6. Take a driving test: Once you feel ready, schedule a driving test with your local department of motor vehicles. If you pass the test, you will receive your driver's license and be able to drive independently.\\n\\nRemember, learning to drive takes time and practice, so be patient with yourself and don't be afraid to ask for help if you need it. Good luck on your journey to becoming a licensed driver!\", error=None, raw_response=\"Learning to drive can be an exciting and rewarding experience. Here are some general steps to help you get started on your journey to becoming a safe and confident driver:\\n\\n1. Get a learner's permit: In most places, you will need to obtain a learner's permit before you can start learning how to drive. Check with your local department of motor vehicles for the specific requirements in your area.\\n\\n2. Take a driver's education course: Consider enrolling in a driver's education course to learn the rules of the road and get some hands-on practice with a qualified instructor.\\n\\n3. Practice with a licensed driver: Before you can get your driver's license, you will need to log a certain number of supervised driving hours with a licensed adult. This is a great opportunity to get comfortable behind the wheel and practice your skills.\\n\\n4. Study the driver's manual: Make sure to familiarize yourself with the driver's manual for your state or country. It contains important information about traffic laws, road signs, and safe driving practices.\\n\\n5. Practice, practice, practice: The more time you spend behind the wheel, the more confident you will become as a driver. Practice in a variety of different conditions and situations to hone your skills.\\n\\n6. Take a driving test: Once you feel ready, schedule a driving test with your local department of motor vehicles. If you pass the test, you will receive your driver's license and be able to drive independently.\\n\\nRemember, learning to drive takes time and practice, so be patient with yourself and don't be afraid to ask for help if you need it. Good luck on your journey to becoming a licensed driver!\")\n",
      "ChatBot: GeneratorOutput(data=\"To get a driver's license in California, you can follow these general steps:\\n\\n1. Obtain a learner's permit: Applicants must be at least 15 and a half years old to apply for a learner's permit in California. You will need to pass a written knowledge test and a vision test to obtain your permit.\\n\\n2. Complete driver's education: If you are under 17 and a half years old, you must complete a driver's education course before applying for a provisional permit.\\n\\n3. Practice driving: With your learner's permit, you can start practicing driving with a licensed adult who is at least 25 years old.\\n\\n4. Apply for a provisional license: After holding your learner's permit for at least 6 months and completing at least 50 hours of practice (including 10 hours at night), you can apply for a provisional license.\\n\\n5. Pass the driving test: Schedule and pass a driving test at a local DMV office. Make sure to bring all required documents and fees.\\n\\n6. Receive your driver's license: If you pass the driving test, you will receive your provisional driver's license. With this license, you will have certain restrictions, such as driving with no passengers under 20 years old for the first year.\\n\\nRemember to check with the California Department of Motor Vehicles (DMV) for the most up-to-date and specific requirements for obtaining a driver's license in the state. Good luck with your journey to becoming a licensed driver in California!\", error=None, raw_response=\"To get a driver's license in California, you can follow these general steps:\\n\\n1. Obtain a learner's permit: Applicants must be at least 15 and a half years old to apply for a learner's permit in California. You will need to pass a written knowledge test and a vision test to obtain your permit.\\n\\n2. Complete driver's education: If you are under 17 and a half years old, you must complete a driver's education course before applying for a provisional permit.\\n\\n3. Practice driving: With your learner's permit, you can start practicing driving with a licensed adult who is at least 25 years old.\\n\\n4. Apply for a provisional license: After holding your learner's permit for at least 6 months and completing at least 50 hours of practice (including 10 hours at night), you can apply for a provisional license.\\n\\n5. Pass the driving test: Schedule and pass a driving test at a local DMV office. Make sure to bring all required documents and fees.\\n\\n6. Receive your driver's license: If you pass the driving test, you will receive your provisional driver's license. With this license, you will have certain restrictions, such as driving with no passengers under 20 years old for the first year.\\n\\nRemember to check with the California Department of Motor Vehicles (DMV) for the most up-to-date and specific requirements for obtaining a driver's license in the state. Good luck with your journey to becoming a licensed driver in California!\")\n",
      "Goodbye!\n"
     ]
    }
   ],
   "source": [
    "chatbot.call()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "lightrag-project",
   "language": "python",
   "name": "light-rag-project"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
